#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
一键创建模型文件
"""

import os
import sys
import subprocess
import shutil
import argparse
import torch


def check_gpu():
    """检查GPU可用性"""
    if torch.cuda.is_available():
        print(f"检测到GPU: {torch.cuda.get_device_name(0)}")
        return True
    else:
        print("未检测到GPU，将使用CPU进行训练")
        return False


def create_model(train=False):
    """一键创建模型文件
    
    Args:
        train (bool): 是否进行训练
    """
    print("开始创建模型文件...")
    
    # 创建model_cache目录
    model_cache_dir = '../model_cache'  # 修改为相对路径
    if not os.path.exists(model_cache_dir):
        os.makedirs(model_cache_dir)
        print(f"创建目录: {model_cache_dir}")
    
    # 检查是否已存在模型文件
    model_dir = os.path.join(model_cache_dir, 'bert_transaction_classifier')
    onnx_model_path = os.path.join(model_cache_dir, 'bert_transaction_classifier.onnx')
    
    if os.path.exists(model_dir) and os.path.exists(onnx_model_path) and not train:
        print("模型文件已存在，无需重新创建。")
        return
    
    # 下载BERT基础模型
    print("下载BERT基础模型...")
    os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
    
    # 检查是否已存在BERT基础模型
    bert_base_dir = os.path.join(model_cache_dir, 'bert-base-chinese')
    if not os.path.exists(bert_base_dir):
        try:
            from transformers import BertTokenizer, BertForSequenceClassification
            tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
            model = BertForSequenceClassification.from_pretrained('bert-base-chinese')
            
            # 保存基础模型
            tokenizer.save_pretrained(bert_base_dir)
            model.save_pretrained(bert_base_dir)
            print(f"BERT基础模型已保存到: {bert_base_dir}")
        except Exception as e:
            print(f"下载BERT基础模型失败: {e}")
            return
    else:
        print(f"BERT基础模型已存在: {bert_base_dir}")
    
    # 如果需要训练，则训练分类模型
    if train:
        # 检查GPU可用性
        use_gpu = check_gpu()
        
        # 训练分类模型
        print("训练分类模型...")
        try:
            # 运行训练脚本，传递GPU参数
            cmd = [sys.executable, 'train_classifier.py']
            if use_gpu:
                cmd.append('--use-gpu')
            
            result = subprocess.run(cmd, 
                                  cwd=os.path.dirname(os.path.abspath(__file__)), 
                                  capture_output=True, 
                                  text=True)
            
            if result.returncode != 0:
                print(f"训练模型失败: {result.stderr}")
                return
            
            print("分类模型训练完成")
        except Exception as e:
            print(f"训练模型时发生错误: {e}")
            return
    else:
        # 不训练，直接下载预训练模型
        print("使用预训练模型...")
        # 这里可以添加下载预训练模型的代码
        pass
    
    # 转换为ONNX格式
    print("转换模型为ONNX格式...")
    try:
        # 运行转换脚本
        result = subprocess.run([sys.executable, 'convert_to_onnx.py'], 
                              cwd=os.path.dirname(os.path.abspath(__file__)), 
                              capture_output=True, 
                              text=True)
        
        if result.returncode != 0:
            print(f"转换模型失败: {result.stderr}")
            return
        
        print("模型已转换为ONNX格式")
    except Exception as e:
        print(f"转换模型时发生错误: {e}")
        return
    
    print("模型文件创建完成！")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='创建模型文件')
    parser.add_argument('--train', action='store_true', help='是否进行训练')
    args = parser.parse_args()
    
    create_model(train=args.train)


if __name__ == "__main__":
    main()